'''
Geometric planes in three dimensions.
Created on Jan 27, 2012

@author: haibo
'''
import numpy as np
from copy import copy

from plane_set import PSF_setd, PSS_flatten_family

def area_volume(planes):
    """ calculate surface area and volume of the space surrounded by the planes
    
    planes is an array of geometry planes (type: Plane).
    Currently I haven't implemented the crystal symmetry. The areas are
    in the same order of the planes. This function also output an
    vertices on each plane, and volume of each plane, in case I want to check 
    the individual divisions.
    This is the volume of the space bounded by the the space angle.
    
    returns [area, volume_total, vertices, volume_space_angle].
    """ 
    n_planes = len(planes)
    # for an enclosed space, there have to be at least 4 planes.
    if n_planes < 4 : return [None, None, None, None]
    
    area = np.zeros(n_planes)
    volume_sa = np.zeros(n_planes)
    vertices = []
    for i in range(n_planes) : vertices.append([])
    
    for i in range(n_planes-2) :
        for j in range(i+1, n_planes-1 ) :
            for k in range(j+1, n_planes) :
                p = planes[i].intercept(planes[j], planes[k])
                if type(p) == np.ndarray :
                    # Trim those extruding points that locate outside 
                    # the enclosed space                    
                    for m in range(n_planes):
                        if m==i or m==j or m==k : continue
                        dist1 = planes[m].distance_signed([0.0, 0, 0])
                        dist2 = planes[m].distance_signed(p)
                        if dist1 * dist2 < -1.0e-15 :  # a small number
                            # it is outside the enclosed space
                            p = None
                            break
                    if type(p) == np.ndarray :
                        vertices[i].append(p)
                        vertices[j].append(p)
                        vertices[k].append(p)
                
    vertices_sorted = []
    for i in range(n_planes) :
        ia_sorted = planes[i].sort_points(vertices[i])
        area[i] = planes[i].area_convex(vertices[i], ia_sorted)
        vertices_sorted.append([vertices[i][x] for x in ia_sorted ])
        # calculate the distance between the planes and origin
        dist = -planes[i].distance_signed([0, 0, 0])
        if dist<0: print "*** warning, inconsistent norm direction."
        volume_sa[i] = 1.0/3 * dist * area[i]
    
    volume_total = sum(volume_sa)
    
    return [volume_total, area, vertices_sorted]


def enclose_polyhedron(number_planes, planes, volume=0.0):
    '''calculates surface-to-volume ratio and surface-families' weights.
    
    plane_family: an array of arrays of planes. Each array contains surfaces 
        of a family.
    volume: a given volume. This parameter is necessary since the
        surface-to-volume depends on volume.
    
    This function can handle one volume. For multiple volumes at constant
    height ratios, eg. for plotting q-d curves, use qf_multivol.
     
    returns: [q, area_weights, h, volume_weights, vertices].
    '''
    volume_calc, area_p, points_p = area_volume(planes)

    # calculate fractional area (or area weight) of each family
    area_ps = []
    index_start = 0
    for subset_n in number_planes:
        index_end = index_start + subset_n
        area_ps.append(np.sum(area_p[index_start : index_end]))
        index_start = index_end
    area_ps = np.array(area_ps)   # convert a list to a numpy.ndarray

    # total surface area
    area_sum = np.sum(area_ps)

    # weight of each surface family
    f_ps = area_ps / area_sum    # array operation: all elements are devided.

    # surface-area-to-volume ratio
    q = area_sum / volume_calc

    # scale the surface areas, point coordinates according to volume
    if volume!=0.0:
        r = np.power(volume / volume_calc, 1.0/3)
        # correct q, f_ps, points_p, (optionally, areal_ps, area_sum)
        q = q / r 
        # f_ps remains the same
        for p in points_p:
            # p is a list of points on one plane
            for xyz in p:
                xyz[0] = xyz[0]*r
                xyz[1] = xyz[1]*r
                xyz[2] = xyz[2]*r

        # area_ps = area_ps * (r*r)  # array operation: all elements are scaled.
        # area_sum = area_sum * (r*r)
        return (q, f_ps, points_p, volume)
    else:
        # r=1.0      # volume is the natural volume_calc.
        return (q, f_ps, points_p, volume_calc)

def _shape_f(dist, ratio_target, number_planes, planes, volume_ref):
    """ function object being used by shape_by_f.
    
    ratio_target: the target fractional surface area.
    number_planes,
    planes,
    volume_ref.
    
    returns: sum of square of difference in ratios of all surface families.
    """
    PSF_setd(number_planes, planes, dist)   # set origin-to-plane distances

    qf = enclose_polyhedron(number_planes, planes, volume_ref)
    #print "target", ratio_target, "real", qf[1], "dist", dist
    t = np.array(ratio_target) - np.array(qf[1])
    t2 = sum([x*x for x in t])
    return t2


def shape_by_f(ratio, lattice, plane_family, volume_ref, distance_initial):
    """ calculates a particle's shape from ratios of surface areas.

    Arguments:
    ratio: a list of ratios of surface areas. They are dimensionless.
        Each ratio corresponds to a plane family, and its value equals to
        the total surface area of all the planes in the plane family divided 
        by the total surface area of all the planes in all the plane families.
    lattice: a CrystalLattice object; it defines the basis and crystal class 
        (or crystallographic point group).
    plane_family: a list of plane families; its length must equal to that of
        ratio.
    volume_ref: the particle's volume.
    Returns:
    a list of [q, f_ps, points_p, volume_calc, distances], where
    q: the surface-area-to-volume ratio.
    f_ps: fractional surface area of each exposed plane.
    points_p: coordinates in each exposed plane.
    distances: origin-to-plane distance of each plane family.
    """
    # import the function fmin_cobyla()
    from scipy.optimize import fmin_cobyla

    # prepare for removing duplicate plane families, e.g., a user may accidentally
    # specify an errornous plane_family = [[1,1,1], [1,0,0], [1,1,1]].
    cc = lattice.crystal_class
    pf_int = cc.plane_family(plane_family[0])   # a list of 3-number tuples
    number_planes = [ len(pf_int) ]
    for p in plane_family[1:]:
        if  not (p in pf_int) :
            pf = cc.plane_family(p)
            pf_int.extend( pf )
            number_planes.append( len(pf) )
        else:
            print("Duplicate plane family found:", p)
    planes = [lattice.geometric_plane(p) for p in pf_int]

    # set params. This is the additional argument that will be passed on to
    # fmin_cobyla()
    params = (ratio, number_planes, planes, volume_ref)
    f = fmin_cobyla(_shape_f, distance_initial, args=params, cons=[], rhobeg=1.0e-1)
    return f
    


# testing and debugging (not exhaustive)
if __name__=="__main__" :
    import crystal_class as CC
    import basis
    from crystal_lattice import CrystalLattice

    diameter = 10.0     # unit: nm
    volume = 4.0/3.0*np.pi* (diameter*0.5)**3

    b1 = basis.cubic(0.4)   # unit: nm
    cc1 = CC.CrystalClass28()
    latt1 = CrystalLattice(b1, cc1)
    number_planes1, planes1 = PSS_flatten_family(
        [[1,2,3], [1,1,0], [2,1,0], [3,3,1], [1,0,0], [1,1,1]], latt1)
    # set distances of the planes
    PSF_setd(number_planes1, planes1, [2.0, 2.0, 2.0, 2.0, 2.0, 2.0])

    qfp = enclose_polyhedron(number_planes1, planes1, volume)
    print "q=", qfp[0]
    print "f=", qfp[1]
    print "numbers of points are", [len(p) for p in qfp[2]]

    latt2 = CrystalLattice(b1, CC.CrystalClass32())
    f = shape_by_f([0.25, 0.75], latt2, [[0,0,1], [1,1,1]], 10.0, [1.0, 1.0])
    print "shape by f", f
    ps = PSS_flatten_family([[0,0,1], [1,1,1]], latt2)
    PSF_setd(ps[0], ps[1], f)
    qfp2 = enclose_polyhedron(ps[0], ps[1], 10.0)
    print "q2=", qfp2[0]
    print "f2=", qfp2[1]

